package com.example.demo.all.controller;


import com.example.demo.all.enums.AiModel;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.*;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.lang.Nullable;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author zwg
 * @date 2025/2/17 15:22
 */
@Slf4j
@Tag(name = "all")
@RestController
@RequestMapping(path = "/all")
public class AllAiController {


    @Autowired
    @Qualifier("openAiChatClient")
    ChatClient openAiChatClient;

    @Autowired
    @Qualifier("zhiPuAiChatClient")
    ChatClient zhiPuAiChatClient;

    @GetMapping("/")
    public String index() {
        return "this is all ai";
    }




    @Operation(summary = "chat")
    @GetMapping("/chat")
    public String testChat(
            @RequestParam AiModel aiModel,
            @RequestParam(defaultValue = "") String system,
            @RequestParam(defaultValue = "") String assistant,
            @RequestParam(defaultValue = "") String user) {

        ChatClient.ChatClientRequestSpec prompt = null;
        switch( aiModel) {
            case OPENAI_CHART -> prompt = openAiChatClient.prompt();
            case ZHIPU_CHAT -> prompt = zhiPuAiChatClient.prompt();
        }

        if(prompt == null){
            return "not find aiModel:"+aiModel;
        }


        log.info("prompt:{}",prompt);
        List<Message> msgs = new ArrayList<>();
        if(StringUtils.isNotBlank(system)){
            msgs.add(new SystemMessage(system));
        }
        if(StringUtils.isNotBlank(assistant)){
            msgs.add(new AssistantMessage(assistant));
        }
        if(StringUtils.isNotBlank(user)){
            msgs.add(new UserMessage(user));
        }
        prompt.messages(msgs);

        try {
            ChatClient.CallResponseSpec response = prompt.call();
            if (response != null) {
                return response.content();
            } else {
                return "not response";
            }
        }catch (Exception e) {
            log.error("",e);
            return e.getMessage();
        }

    }

    /*@Operation(summary = "chatStreamResp")
    @GetMapping("/chatStreamResp")
    public String chatStreamResp(
            @RequestParam AiModel aiModel,
            @RequestParam(defaultValue = "") String system,
            @RequestParam(defaultValue = "") String assistant,
            @RequestParam(defaultValue = "") String user) {

        ChatClient.ChatClientRequestSpec prompt = null;
        switch( aiModel) {
            case OPENAI_CHART -> prompt = openAiChatClient.prompt();
            case ZHIPU_CHAT -> prompt = zhiPuAiChatClient.prompt();
        }

        try {
            if(prompt == null){
                return "not find aiModel:"+aiModel;
            }

            List<Message> msgs = new ArrayList<>();
            if(StringUtils.isNotBlank(system)){
                msgs.add(new SystemMessage(system));
            }
            if(StringUtils.isNotBlank(assistant)){
                msgs.add(new AssistantMessage(assistant));
            }
            if(StringUtils.isNotBlank(user)){
                msgs.add(new UserMessage(user));
            }
            prompt.messages(msgs);

            ChatClient.StreamResponseSpec response = prompt.stream();
            log.info("resp:=====================");
            AtomicInteger count = new AtomicInteger();
            Flux<ChatResponse> responseFlux = response.chatResponse();
            responseFlux.buffer(50).subscribe((chatResponseList)->{
                StringBuilder message = new StringBuilder();
                for (ChatResponse chatResponse : chatResponseList) {
                    message.append(getContentFromChatResponse(chatResponse));
                }
                log.info("resp:{} {}", count.getAndIncrement(), message.toString());
            });
        }catch (Exception e) {
            log.error("",e);
        }

        return "";
    }*/


    private static String getContentFromChatResponse(@Nullable ChatResponse chatResponse) {
        return Optional.ofNullable(chatResponse)
                .map(ChatResponse::getResult)
                .map(Generation::getOutput)
                .map(AbstractMessage::getText)
                .orElse(null);
    }


}
